package edu.gatech.cc.liam.core.rl.models.decpomdp;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.Map.Entry;

import edu.gatech.cc.liam.core.Globals;
import edu.gatech.cc.liam.core.rl.models.CartesianIterator;
import edu.gatech.cc.liam.core.rl.models.JointIterator;
import edu.gatech.cc.liam.geometry.linear.FeasibleSet;
import edu.gatech.cc.liam.geometry.linear.Halfspace;
import edu.gatech.cc.liam.geometry.linear.NPoint;
import edu.gatech.cc.liam.geometry.linear.SparsePoint;
import edu.gatech.cc.liam.geometry.linear.Vector;
import edu.gatech.cc.liam.geometry.linear.convexhull.PolytopeProjectionPipe;
import edu.gatech.cc.liam.marl.decpomdp.Belief;

public class DecMDP {
	public final double discount;
	public final int numAgents;
	public final List<String>[] types;  // a.k.a states
	public final List<String>[] actions;
//	private ProbabilityMap<Pair<List<String>, List<String>>, List<String>> transitionFunction;
//	private RewardFunction<Pair<List<String>, List<String>>> rewardFunction;
	private Vector[][] transitionFunction; // [joint-type-start][joint-action][joint-type-end]
	private double[][] rewardFunction; // reward of [joint-type][joint-action]
	private double[] startDistr;
	
	// constructed fields
	public int numJTypes;
	public int numJActions;
	public boolean[] validTypeList;
//	public List<List<String>> allJointActions;
//	public List<List<String>> allJointTypes;
	public int[] typeSS;
	public int[] actionSS;
	
	public DecMDP noNature = null;
	private int[] noNtrIndexTransform;
	
	public DecMDP(List<String>[] types, List<String>[] actions, double discount) {
		assert(actions.length == types.length);
		this.discount = discount;
		this.types = types;
		this.actions = actions;
		this.numAgents = actions.length;
		numJTypes = 1;
		numJActions = 1;
		typeSS = new int[numAgents];
		actionSS = new int[numAgents];
		for(int i=0; i<numAgents; i++){
			typeSS[i] = types[i].size();
			numJTypes *= typeSS[i];
			actionSS[i] = actions[i].size();
			numJActions *= actionSS[i];
		}
//		allJointTypes = makeJointList(types, typeSS);
//		allJointActions = makeJointList(actions, actionSS);
	}
	
	public DecMDP(List<String>[] types,
				List<String>[] actions,
				Vector[][] transitionFunction,
				double[][] rewardFunction,
				double[] startDistr,
				double discount) {
		this(types, actions, discount);
		setFunctions(transitionFunction, rewardFunction, startDistr);
	}

	public void makeValidTypeList() {
		this.validTypeList = new boolean[numJTypes];
		for(int jTi = 0; jTi < numJTypes; jTi++) {
			validTypeList[jTi] = true;
			if (! transitionFunction[jTi][0].isDistribution())
				validTypeList[jTi] = false;
		}
	}
	public boolean isValidType(int jTi) {
		return validTypeList[jTi];
	}
	
	public static DecMDP makeRandomDecMDP(int numTypes, int numActions, int numAgents) {
		ArrayList<String> typeSpace = new ArrayList<String>();
		for(int i=0; i<numTypes; i++) {
			typeSpace.add(Integer.toString(i));
		}
		ArrayList<String> actionSpace = new ArrayList<String>();
		for(int i=0; i<numActions; i++) {
			actionSpace.add(Integer.toString(i));
		}
		List<String>[] types = new ArrayList[numAgents];
		List<String>[] actions = new ArrayList[numAgents];
		for(int i=0; i<numAgents; i++) {
			types[i] = typeSpace;
			actions[i] = actionSpace;
		}
		return new DecMDP(types, actions, 0.9);
	}
	
	private void makeNoNatureMDP() {
		// combine types of agents with only 1 action (e.g. nature)
		List<String>[] modTypes = Arrays.copyOf(types, numAgents);
		for(int i=0; i<numAgents; i++) {
			if(actions[i].size() == 1) {
				modTypes[i] = new ArrayList<String>();
				modTypes[i].add("combined");
			}
		}
		noNature = new DecMDP(modTypes, actions, Double.NaN);
		// make the index transform
		noNtrIndexTransform = new int[numJTypes * numJActions];
		for(int[] jT : new CartesianIterator(typeSS)) {
			int jTi = getJointTypeIndex(jT);
			for(int i=0; i<numAgents; i++) {
				if(actions[i].size() == 1) {
					jT[i] = 0;
				}
			}
			int noNatureJTi = this.noNature.getJointTypeIndex(jT);
			for(int[] jA : new CartesianIterator(actionSS)) {
				int jAi = getJointActionIndex(jA);
				int fullIndex = this.getVarIndex(jTi, jAi);
				int noNatureIndex = this.noNature.getVarIndex(noNatureJTi, jAi);
				noNtrIndexTransform[fullIndex] = noNatureIndex;
			}
		}
	}
	
	public void setFunctions(Vector[][] transitionFunction, double[][] rewardFunction,
				        	 double[] startDistr) {
		assert this.transitionFunction == null;
		assert this.rewardFunction == null;
		assert this.startDistr == null;
		this.transitionFunction = transitionFunction;
		this.rewardFunction = rewardFunction;
		this.startDistr = startDistr;
		// if functions are set, then this is a true model (not a no-nature model) so make the
		// no nature model.
		makeNoNatureMDP();
		makeValidTypeList();
	}
	
	public static List<List<String>> makeJointList(List<String>[] sets, int[] setSizes) {
		List<List<String>> collector = new ArrayList<List<String>>();
		assert(sets.length == setSizes.length);
		for(int[] ji : new CartesianIterator(setSizes)) {
		    String[] j = new String[setSizes.length];
			for(int i=0; i<setSizes.length; i++) {
				j[i] = sets[i].get(ji[i]);
			}
			collector.add(Arrays.asList(j));
		}
		return collector;
	}
	
	public int getJointTypeIndex(int[] typeInds) {
		return JointIterator.getJointIndex(typeSS, typeInds);
	}
	public int getJointActionIndex(int[] actionInds) {
		return JointIterator.getJointIndex(actionSS, actionInds);	
	}

	/**
	 * @param jType the initial state
	 * @param jActions the joint actions taken
	 * @return a probability distribution over successor states s.t. return[sJT] is the probability
	 * of ending up at successor state sJT after taking jActions in jType
	 */
	public Vector getTransitionProbs(int jType, int jActions) {
		return transitionFunction[jType][jActions];
	}
	public double getReward(int jTypeOld, int jActions) {
		return rewardFunction[jTypeOld][jActions];
	}
	public double[] getStartDist() {
		return startDistr;
	}
	
	public Belief getSuccessorBelief(Belief b, DecMDPSingleStagePolicy ssp) {
		SparsePoint newB = new SparsePoint(b.dimensions());
		for(Entry<Integer, Double> e : b.getProbDist()) {
			int jTi = e.getKey();
			assert(isValidType(jTi));
			for(int jAi=0; jAi < numJActions; jAi++) {
				double actionProb = ssp.getProb(jTi,jAi);
				if(actionProb > 0.0) {
					Vector JTAfterJTJA = getTransitionProbs(jTi, jAi).clone();
					assert JTAfterJTJA.isDistribution();
					JTAfterJTJA.scale(actionProb);
					JTAfterJTJA.scale(e.getValue());
					newB.add(JTAfterJTJA);
				}
			}
		}
		return new Belief(newB);
	}

	public Set<Belief> sampleBeliefs(Belief startB, DecMDPPolicy policy,
									 int numToSample, int depth, boolean verbose) {
		HashSet<Belief> bSet = new HashSet<Belief>();
		final int MAX_SAMPLE_TRIES = numToSample * 10;
		int numTries = 0;
		while (bSet.size() < numToSample && numTries < MAX_SAMPLE_TRIES) {
			numTries++;
			StringBuffer vBuf = verbose?new StringBuffer():null;
			ArrayList<Belief> trace = sampleBeliefTrace(startB, policy, depth, vBuf);
			Belief newB = trace.get(trace.size()-1);
			if (!Belief.softContains(bSet, newB)) {
				if (vBuf!=null)  System.out.println(vBuf.toString());
				bSet.add(newB);
			}
		}
		return bSet;
	}

	public ArrayList<Belief> sampleBeliefTrace(Belief startB, DecMDPPolicy policy,
								     			int depth, StringBuffer vBuf) {
		ArrayList<Belief> trace = new ArrayList<Belief>();
		Belief newB = startB.clone();
		if (vBuf!=null) vBuf.append("\nStarting at " + newB);
		for (int i = 0; i < depth; i++) {
			newB.validate(this.validTypeList);
			DecMDPSingleStagePolicy ssp = policy.getActions(newB);
			assert ssp != null;
			if (vBuf!=null) vBuf.append("\ntook action " + ssp.getActionString(newB));
			newB = getSuccessorBelief(newB, ssp);
			if (vBuf!=null) vBuf.append("got to " + newB);
			trace.add(newB);
		}
		return trace;
	}

	
	/**
	 * @return a matrix s.t. return[joint-type-index][joint-action-index] = prob of taking that
	 * joint-action in the joint-type. This should be random s.t. only one joint-action for every
	 * joint-type is taken. Each type takes an action and the product of the type's actions is given
	 * as the joint-action.
	 */
	public DecMDPPolicy getRandomPolicy() {
		class RandomActionGen implements DecMDPPolicy {
			public DecMDPSingleStagePolicy getActions(Belief publicB) {
				return makeRandomPolicy();
			}
		}
		return new RandomActionGen();
	}
	/**
	 * for each player and type select an action randomly
	 */
	public DecMDPSingleStagePolicy makeRandomPolicy() {
		int[][] typeAtions =  getRandomActionForEachType();
		return new DecMDPSingleStagePolicy(this, typeAtions);
}
	
	public DecMDPPolicy getSymRandomPolicy() {
		class RandomSymActionGen implements DecMDPPolicy {
			// for each player and type select an action randomly
			public DecMDPSingleStagePolicy getActions(Belief publicB) {
				// for each type select an action randomly, give to all players (except nature)
				int[][] typeActions = new int[numAgents][];
				int[] as = new int[typeSS[1]];
				for(int t=0; t<typeSS[1]; t++) {
					as[t] = Globals.rand.nextInt(actionSS[1]);
				}
				for(int agent=0; agent<numAgents; agent++) {
					if(actionSS[agent] > 1) {
						assert typeSS[agent] == typeSS[1];
						assert actionSS[agent] == actionSS[1];
						typeActions[agent] = as;
					} else {
						typeActions[agent] = new int[typeSS[agent]];
					}
				}
				return new DecMDPSingleStagePolicy(DecMDP.this, typeActions);
			}
		}
		return new RandomSymActionGen();
	}
	
	private int[][] getRandomActionForEachType() {
		int[][] typeAtions = new int[numAgents][];
		for(int agent=0; agent<numAgents; agent++) {
			typeAtions[agent] = new int[typeSS[agent]];
			for(int t=0; t<typeSS[agent]; t++) {
				typeAtions[agent][t] = Globals.rand.nextInt(actionSS[agent]);
			}
		}
		return typeAtions;
	}

	public int[] getFirstPossibleJType(Belief b) {
		int[] firstPossibleJT = null;
		for(int[] typeInds : new CartesianIterator(typeSS)) {
			int jTi = getJointTypeIndex(typeInds);
			if(b.getProb(jTi) > Globals.TOLERANCE) {
				firstPossibleJT = typeInds;
				break;
			}
		}
		return firstPossibleJT;
	}
	
	/*
	 * @return matrix of boolean s.t. return[agent][type-index] = true if there is positive
	 * probability that the agent has type 'type-index'
	 */
	public boolean[][] whichTypesInBelief(Belief b) {
		boolean[][] typePossible = new boolean[numAgents][]; 
		for(int agent = 0; agent < numAgents; agent++)
			typePossible[agent] = new boolean[typeSS[agent]]; // default is false
		for(int[] typeInds : new CartesianIterator(typeSS)) {
			int jTi = getJointTypeIndex(typeInds);
			for(int agent = 0; agent < numAgents; agent++) {
				if(b.getProb(jTi) > Globals.TOLERANCE)
					typePossible[agent][typeInds[agent]] = true;
			}
		}
		return typePossible;
	}
	
	public FeasibleSet makePolicyFeasibleSet() {
		// define linear program over variables t_a
		// one for each (t = joint type, a = joint action)
		int numVars = numJTypes *  this.numJActions;
		FeasibleSet policyFS = new FeasibleSet();
		// All variables range from 0 to 1.  Each represents the probability that
		// joint-action a will be taken if the agents have joint-type t.
		policyFS.lb = new NPoint(numVars, 0.0);
		policyFS.ub = new NPoint(numVars, 1.0);
//		for(int i=0; i<numVars; i++) {
//			NPoint normal = new NPoint(numVars, 0.0);
//			normal.values[i] = -1.0;
//			policyFS.add(new Halfspace(normal, 0.0));
//			normal.values[i] = 1.0;
//			policyFS.add(new Halfspace(normal, -1.0));
//		}
		// for every joint-type, all joint-actions must sum to 1 (a probability distribution)
		for(int i=0; i < this.numJTypes; i++) {
			NPoint constraint  = new NPoint(numVars, 0.0);
			for(int j=0; j < this.numJActions; j++) {
//				System.out.println(this.getVarIndex(i, j) + ": jT" + i + " jA" + j);
				constraint.values[this.getVarIndex(i, j)] = 1.0;
			}
			policyFS.equalities.add(new Halfspace(constraint, -1.0));
		}
		// sum over actions is equal over types
		// for every agent, their type should not change the probability of other agent's 
		// joint-actions. In other words, if an agent had only one action
		// (by summing over joint-actions fixing that agent's action) then the joint-action across
		// all types of that agent should be equal.
		// This insures that the policy doesn't assume information is shared across agents
		for(int agent = 0; agent < this.numAgents; agent++) { // for each agent i
			int[] numT = this.typeSS.clone();
			numT[agent] = 1; 									// don't iterate over i's types
			for(int[] typeInds : new CartesianIterator(numT)) { // for each joint-type_{-i)
				int[] numA = this.actionSS.clone();
				numA[agent] = 1;								// don't iterate over i's actions
				for(int[] actionInds : new CartesianIterator(numA)) { // for each joint-action_{-i}
					// create lists (one for each agent i type) of the joint-actions
					// over the agent's possible actions (flattening jTjA in agent i's direction).
					// The other agent's actions and types are specified at this point.
					ArrayList<NPoint> sumOfActionsForEachJointType = new ArrayList<NPoint>();
					for(int typeNum = 0; typeNum < this.typeSS[agent]; typeNum++) {	
						typeInds[agent] = typeNum;
						int jointTypeIndex = this.getJointTypeIndex(typeInds);
						NPoint actionSum = new NPoint(numVars, 0.0);
						for(int actionNum = 0; actionNum < this.actionSS[agent]; actionNum++) {
							actionInds[agent] = actionNum;
							int jointActionIndex = this.getJointActionIndex(actionInds);
							actionSum.values[this.getVarIndex(jointTypeIndex, jointActionIndex)] = 1.0;
						}
						sumOfActionsForEachJointType.add(actionSum);
					}
					// now create sumOfActionsForEachJointType.size() - 1 equality constraints
					// (skip the first) such that the sum of each is equal to the first
					for(int i=1; i<sumOfActionsForEachJointType.size(); i++) {
						NPoint eqConstraint = new NPoint(numVars, 0.0);
						eqConstraint.subtract(sumOfActionsForEachJointType.get(0));
						eqConstraint.add(sumOfActionsForEachJointType.get(i));
						policyFS.equalities.add(new Halfspace(eqConstraint, 0.0));
					}
				}
			}
		}
		return policyFS;
	}
	
	public int getVarIndex(int jTypeIndex, int jActionIndex) {
		assert jTypeIndex < numJTypes;
		assert jActionIndex < numJActions;
		return (numJActions * jTypeIndex) + jActionIndex;
	}
	public int getVarIndex(int[] jType, int[] jAction) {
		return getVarIndex(getJointTypeIndex(jType), getJointActionIndex(jAction));
	}
	public int getNoNatureIndex(int varIndex) {
		return noNtrIndexTransform[varIndex];
	}
	
	public static void main(String[] args) {
		int maxActions = 8;
		int maxTypes = 8;
		int maxAgents = 2;
		System.out.println("numActions, numTypes, numAgents, numVars, numEq, numInEQ, TrueDims");
		for(int numActions = 1; numActions <= maxActions; numActions++) {
			for(int numTypes = 1; numTypes <= maxTypes; numTypes++) {
				for(int numAgents = 1; numAgents <= maxAgents; numAgents++) {
					DecMDP m = DecMDP.makeRandomDecMDP(numTypes, numActions, numAgents);
					FeasibleSet fs = m.makePolicyFeasibleSet();
					System.out.println(numActions + ", " +
							numTypes + ", " +
							numAgents + ", " +
							fs.getDimensions() + ", " +
							fs.equalities.size() + ", " +
							fs.inequalities.size() + ", " +
							PolytopeProjectionPipe.projectToManifold(fs.transformToPolytope()).getDimensions()
							);
				}
			}
		}
	}
}
